package tools

import (
	"os"
	"testing"
)

// ExitMock registers Exit (os.Exit) calls
type ExitMock struct {
	DefaultCode int

	lastCodes []int
}

// LastCode returns exit code for the last Exit call event.
func (e *ExitMock) LastCode() int {
	if len(e.lastCodes) > 0 {
		return e.lastCodes[len(e.lastCodes)-1]
	}
	return e.DefaultCode
}

func (e *ExitMock) exit(code int) {
	e.lastCodes = append(e.lastCodes, code)
}

// WithExit sets exit mock for testing purposes.
func WithExit(t testing.TB, mock *ExitMock) {
	origin := Exit
	Exit = mock.exit

	t.Cleanup(func() {
		Exit = origin
	})
}

// WithArgs sets application arguments.
func WithArgs(t testing.TB, args []string) {
	orig := os.Args
	os.Args = args

	t.Cleanup(func() {
		os.Args = orig
	})
}
